"""
此脚本用于连接到MCP服务器并将输入和输出传输到websocket端点。
版本: 0.1.0

用法:

# 在config/mcp_config.py中设置MCP_ENDPOINT
python mcp_pipe.py <mcp_script>

"""

import asyncio
import websockets
import subprocess
import logging
import os
import signal
import sys
import random
from dotenv import load_dotenv

# 添加当前目录到Python路径
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# 从.env文件加载环境变量
load_dotenv()

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('MCP_PIPE')

# 重连设置
INITIAL_BACKOFF = 1  # 初始等待时间（秒）
MAX_BACKOFF = 600  # 最大等待时间（秒）
reconnect_attempt = 0
backoff = INITIAL_BACKOFF

# 导入MCP配置
try:
    from config.mcp_config import MCP_CONFIG
except ImportError as e:
    logger.error(f"无法导入config/mcp_config.py配置文件: {e}")
    logger.error("请确保config/mcp_config.py文件存在并包含MCP_CONFIG字典")
    sys.exit(1)

async def connect_with_retry(uri):
    """使用重试机制连接到WebSocket服务器"""
    global reconnect_attempt, backoff
    while True:  # 无限重连
        try:
            if reconnect_attempt > 0:
                wait_time = backoff * (1 + random.random() * 0.1)  # 添加一些随机抖动
                logger.info(f"等待 {wait_time:.2f} 秒后进行第 {reconnect_attempt} 次重连尝试...")
                await asyncio.sleep(wait_time)
                
            # 尝试连接
            await connect_to_server(uri)
        
        except Exception as e:
            reconnect_attempt += 1
            logger.warning(f"连接关闭 (尝试: {reconnect_attempt}): {e}")            
            # 计算下次重连的等待时间（指数退避）
            backoff = min(backoff * 2, MAX_BACKOFF)

async def connect_to_server(uri):
    """连接到WebSocket服务器并与`mcp_script`建立双向通信"""
    global reconnect_attempt, backoff
    process = None  # 初始化process变量
    try:
        logger.info(f"正在连接到WebSocket服务器...")
        async with websockets.connect(uri) as websocket:
            logger.info(f"成功连接到WebSocket服务器")
            
            # 如果连接正常关闭，则重置重连计数器
            reconnect_attempt = 0
            backoff = INITIAL_BACKOFF
            
            # 启动mcp_script进程
            process = subprocess.Popen(
                [sys.executable, mcp_script],
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                encoding='utf-8',
                text=True  # 使用文本模式
            )
            logger.info(f"已启动 {mcp_script} 进程")
            
            # 创建两个任务：从WebSocket读取并写入进程，从进程读取并写入WebSocket
            await asyncio.gather(
                pipe_websocket_to_process(websocket, process),
                pipe_process_to_websocket(process, websocket),
                pipe_process_stderr_to_terminal(process)
            )
    except websockets.exceptions.ConnectionClosed as e:
        logger.error(f"WebSocket连接关闭: {e}")
        raise  # 重新抛出异常以触发重连
    except Exception as e:
        logger.error(f"连接错误: {e}")
        raise  # 重新抛出异常
    finally:
        # 确保子进程被正确终止
        if process is not None:
            logger.info(f"正在终止 {mcp_script} 进程")
            try:
                process.terminate()
                process.wait(timeout=5)
            except subprocess.TimeoutExpired:
                process.kill()
            except Exception as e:
                logger.error(f"终止进程时出错: {e}")
            logger.info(f"{mcp_script} 进程已终止")

async def pipe_websocket_to_process(websocket, process):
    """从WebSocket读取数据并写入进程stdin"""
    try:
        while True:
            # 从WebSocket读取消息
            message = await websocket.recv()
            logger.debug(f"<< {message[:120]}...")
            
            # 写入进程stdin（在文本模式下）
            if isinstance(message, bytes):
                message = message.decode('utf-8')
            process.stdin.write(message + '\n')
            process.stdin.flush()
    except Exception as e:
        logger.error(f"WebSocket到进程管道错误: {e}")
        raise  # 重新抛出异常以触发重连
    finally:
        # 关闭进程stdin
        if not process.stdin.closed:
            process.stdin.close()

async def pipe_process_to_websocket(process, websocket):
    """从进程stdout读取数据并发送到WebSocket"""
    try:
        while True:
            # 从进程stdout读取数据
            data = await asyncio.get_event_loop().run_in_executor(
                None, process.stdout.readline
            )
            
            if not data:  # 如果没有数据，进程可能已结束
                logger.info("进程已结束输出")
                break
                
            # 发送数据到WebSocket
            logger.debug(f">> {data[:120]}...")
            # 在文本模式下，数据已经是字符串，无需解码
            await websocket.send(data)
    except Exception as e:
        logger.error(f"进程到WebSocket管道错误: {e}")
        raise  # 重新抛出异常以触发重连

async def pipe_process_stderr_to_terminal(process):
    """从进程stderr读取数据并打印到终端"""
    try:
        while True:
            # 从进程stderr读取数据
            data = await asyncio.get_event_loop().run_in_executor(
                None, process.stderr.readline
            )
            
            if not data:  # 如果没有数据，进程可能已结束
                logger.info("进程已结束stderr输出")
                break
                
            # 将stderr数据打印到终端（在文本模式下，数据已经是字符串）
            sys.stderr.write(data)
            sys.stderr.flush()
    except Exception as e:
        logger.error(f"进程stderr管道错误: {e}")
        raise  # 重新抛出异常以触发重连

def signal_handler(sig, frame):
    """处理中断信号"""
    logger.info("收到中断信号，正在关闭...")
    sys.exit(0)

if __name__ == "__main__":
    # 注册信号处理程序
    signal.signal(signal.SIGINT, signal_handler)
    
    # mcp_script
    if len(sys.argv) < 2:
        logger.error("用法: mcp_pipe.py <mcp_script>")
        sys.exit(1)
    
    mcp_script = sys.argv[1]
    
    # 从配置文件获取令牌
    endpoint_url = MCP_CONFIG.get('MCP_ENDPOINT')
    if not endpoint_url:
        logger.error("请在config/mcp_config.py中设置MCP_ENDPOINT")
        sys.exit(1)
    
    # 启动主循环
    try:
        asyncio.run(connect_with_retry(endpoint_url))
    except KeyboardInterrupt:
        logger.info("程序被用户中断")
    except Exception as e:
        logger.error(f"程序执行错误: {e}")